#include "common.h"
#include <libgen.h> // 添加basename函数的头文件

// 文件接收状态结构体
typedef struct {
    int fd;                     // 文件描述符
    char filename[256];         // 文件名
    size_t filesize;            // 文件大小
    size_t received_size;       // 已接收大小
    size_t last_block;          // 最后接收的块索引
    int block_size;             // 块大小
} FileReceiveState;

// 线程参数结构体
typedef struct {
    int client_fd;              // 客户端socket
    Queue* data_queue;          // 数据队列
} ThreadArgs;

// 保存文件接收状态到文件，用于断点续传
void save_resume_state(FileReceiveState* state) {
    char resume_filename[270];
    snprintf(resume_filename, sizeof(resume_filename), "%s.resume", state->filename);
    
    FILE* fp = fopen(resume_filename, "wb");
    if (!fp) {
        perror("Failed to save resume state");
        return;
    }
    
    fwrite(state, sizeof(FileReceiveState), 1, fp);
    fclose(fp);
}

// 从文件加载接收状态，用于断点续传
int load_resume_state(const char* filename, FileReceiveState* state) {
    char resume_filename[270];
    snprintf(resume_filename, sizeof(resume_filename), "%s.resume", filename);
    
    FILE* fp = fopen(resume_filename, "rb");
    if (!fp) {
        return -1;
    }
    
    size_t read_size = fread(state, sizeof(FileReceiveState), 1, fp);
    fclose(fp);
    
    return (read_size == 1) ? 0 : -1;
}

// 删除断点续传状态文件
void delete_resume_state(const char* filename) {
    char resume_filename[270];
    snprintf(resume_filename, sizeof(resume_filename), "%s.resume", filename);
    unlink(resume_filename);
}

// 写入线程函数
void* write_thread(void* arg) {
    Queue* data_queue = (Queue*)arg;
    FileReceiveState state = {0};
    int file_opened = 0;
    
    while (1) {
        size_t size;
        void* data = queue_pop(data_queue, &size);
        if (!data) continue;
        
        // 处理文件信息包
        if (size == sizeof(FileInfoPacket)) {
            FileInfoPacket* info = (FileInfoPacket*)data;
            if (info->type == FILE_INFO) {
                // 检查是否存在断点续传状态
                if (load_resume_state(info->filename, &state) == 0) {
                    printf("Resuming file transfer for %s from block %zu\n", 
                           info->filename, state.last_block + 1);
                    
                    // 打开文件进行追加
                    state.fd = open(info->filename, O_WRONLY | O_CREAT, 0644);
                    if (state.fd < 0) {
                        perror("Failed to open file for resuming");
                        free(data);
                        continue;
                    }
                    
                    // 定位到断点位置
                    if (lseek(state.fd, state.received_size, SEEK_SET) < 0) {
                        perror("Failed to seek to resume position");
                        close(state.fd);
                        free(data);
                        continue;
                    }
                    
                    file_opened = 1;
                } else {
                    // 新文件传输
                    strncpy(state.filename, info->filename, sizeof(state.filename) - 1);
                    state.filesize = info->filesize;
                    state.received_size = 0;
                    state.last_block = -1;
                    state.block_size = info->block_size;
                    
                    // 创建新文件
                    state.fd = open(info->filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
                    if (state.fd < 0) {
                        perror("Failed to create file");
                        free(data);
                        continue;
                    }
                    
                    file_opened = 1;
                    printf("Starting to receive file: %s (size: %zu bytes)\n", 
                           info->filename, info->filesize);
                }
            }
        }
        // 处理文件数据包
        else if (size == sizeof(FileDataPacket) && file_opened) {
            FileDataPacket* packet = (FileDataPacket*)data;
            if (packet->type == FILE_DATA) {
                // 检查是否为期望的下一个块
                if (packet->block_index == state.last_block + 1) {
                    ssize_t written = write(state.fd, packet->data, packet->data_size);
                    if (written > 0) {
                        state.received_size += written;
                        state.last_block = packet->block_index;
                        
                        // 每接收10个块保存一次状态
                        if (packet->block_index % 10 == 0) {
                            save_resume_state(&state);
                            printf("Received block %zu (%.2f%%)\n", 
                                   packet->block_index, 
                                   (float)state.received_size / state.filesize * 100);
                        }
                    }
                }
            }
        }
        // 处理文件结束包
        else if (size == sizeof(FileEndPacket) && file_opened) {
            FileEndPacket* end = (FileEndPacket*)data;
            if (end->type == FILE_END) {
                printf("File transfer completed: %s (%zu bytes)\n", 
                       state.filename, state.received_size);
                
                if (state.fd > 0) {
                    close(state.fd);
                    state.fd = -1;
                }
                
                // 传输完成，删除断点续传状态
                delete_resume_state(state.filename);
                file_opened = 0;
            }
        }
        
        free(data);
    }
    
    return NULL;
}

// 接收线程函数
void* receive_thread(void* arg) {
    ThreadArgs* args = (ThreadArgs*)arg;
    int client_fd = args->client_fd;
    Queue* data_queue = args->data_queue;
    
    while (1) {
        // 先接收包类型
        PacketType type;
        ssize_t recv_size = recv(client_fd, &type, sizeof(PacketType), 0);
        
        if (recv_size <= 0) {
            if (recv_size == 0) {
                printf("Client disconnected\n");
            } else {
                perror("recv failed");
            }
            break;
        }
        
        // 根据包类型接收不同的数据
        switch (type) {
            case FILE_INFO: {
                FileInfoPacket info;
                info.type = type;
                
                recv_size = recv(client_fd, ((char*)&info) + sizeof(PacketType), 
                                sizeof(FileInfoPacket) - sizeof(PacketType), 0);
                if (recv_size > 0) {
                    queue_push(data_queue, &info, sizeof(FileInfoPacket));
                    
                    // 检查是否需要断点续传
                    FileReceiveState state;
                    if (load_resume_state(info.filename, &state) == 0) {
                        // 发送断点续传请求
                        FileResumePacket resume;
                        resume.type = FILE_RESUME;
                        strncpy(resume.filename, info.filename, sizeof(resume.filename) - 1);
                        resume.last_block = state.last_block;
                        
                        send(client_fd, &resume, sizeof(FileResumePacket), 0);
                    }
                }
                break;
            }
            case FILE_DATA: {
                FileDataPacket packet;
                packet.type = type;
                
                // 接收块索引和数据大小
                recv_size = recv(client_fd, ((char*)&packet) + sizeof(PacketType), 
                                sizeof(size_t) * 2, 0);
                if (recv_size > 0) {
                    // 接收数据内容
                    recv_size = recv(client_fd, packet.data, packet.data_size, 0);
                    if (recv_size > 0) {
                        queue_push(data_queue, &packet, sizeof(FileDataPacket));
                    }
                }
                break;
            }
            case FILE_END: {
                FileEndPacket end;
                end.type = type;
                
                recv_size = recv(client_fd, ((char*)&end) + sizeof(PacketType), 
                                sizeof(FileEndPacket) - sizeof(PacketType), 0);
                if (recv_size > 0) {
                    queue_push(data_queue, &end, sizeof(FileEndPacket));
                }
                break;
            }
            default:
                printf("Unknown packet type: %d\n", type);
                break;
        }
    }
    
    close(client_fd);
    free(args);
    return NULL;
}

int main() {
    int server_fd, client_fd;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);
    
    // 创建socket
    server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) {
        HANDLE_ERROR("socket failed");
    }
    
    // 设置地址重用
    int opt = 1;
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        HANDLE_ERROR("setsockopt failed");
    }
    
    // 绑定地址和端口
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(SERVER_PORT);
    
    if (bind(server_fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        HANDLE_ERROR("bind failed");
    }
    
    // 监听连接
    if (listen(server_fd, 5) < 0) {
        HANDLE_ERROR("listen failed");
    }
    
    printf("Server started, listening on port %d\n", SERVER_PORT);
    
    // 创建数据队列
    Queue* data_queue = queue_init();
    
    // 创建写入线程
    pthread_t write_tid;
    if (pthread_create(&write_tid, NULL, write_thread, data_queue) != 0) {
        HANDLE_ERROR("Failed to create write thread");
    }
    
    // 主循环，接受客户端连接
    while (1) {
        client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &client_len);
        if (client_fd < 0) {
            perror("accept failed");
            continue;
        }
        
        printf("New client connected: %s:%d\n", 
               inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
        
        // 创建接收线程参数
        ThreadArgs* args = (ThreadArgs*)malloc(sizeof(ThreadArgs));
        if (!args) {
            perror("Failed to allocate thread args");
            close(client_fd);
            continue;
        }
        
        args->client_fd = client_fd;
        args->data_queue = data_queue;
        
        // 创建接收线程
        pthread_t recv_tid;
        if (pthread_create(&recv_tid, NULL, receive_thread, args) != 0) {
            perror("Failed to create receive thread");
            free(args);
            close(client_fd);
            continue;
        }
        
        // 分离线程，让其自行结束
        pthread_detach(recv_tid);
    }
    
    // 清理资源
    close(server_fd);
    queue_destroy(data_queue);
    
    return 0;
} 